import librosa
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
import pandas as pd
import os
# 音频加载与预处理
def load_audio(filepath, target_sr=16000, max_duration=10):
    """
    加载音频文件并进行标准化处理
    :param filepath: 音频文件路径
    :param target_sr: 目标采样率
    :param max_duration: 最大音频长度（秒）
    """
    audio, sr = librosa.load(filepath, sr=target_sr)
    if len(audio) > target_sr * max_duration:
        audio = audio[:target_sr * max_duration]
    else:
        padding = target_sr * max_duration - len(audio)
        audio = np.pad(audio, (0, padding))
    return audio

def extract_mel_spectrogram(audio, sr=16000, n_mels=128, n_fft=2048, hop_length=512):
    """
    提取梅尔频谱特征
    :param audio: 音频信号
    :param sr: 采样率
    :param n_mels: 梅尔频谱数量
    """
    mel_spec = librosa.feature.melspectrogram(y=audio, sr=sr, n_mels=n_mels,
                                              n_fft=n_fft, hop_length=hop_length)
    log_mel_spec = librosa.power_to_db(mel_spec, ref=np.max)
    return log_mel_spec

# 自定义数据集
class BirdDataset(Dataset):
    def __init__(self, file_paths, labels, sr=16000, max_duration=10, n_mels=128):
        """
        数据集初始化
        :param file_paths: 音频文件路径列表
        :param labels: 标签列表
        """
        self.file_paths = file_paths
        self.labels = labels
        self.sr = sr
        self.max_duration = max_duration
        self.n_mels = n_mels

    def __len__(self):
        return len(self.file_paths)

    def __getitem__(self, idx):
        filepath = self.file_paths[idx]
        label = self.labels[idx]
        audio = load_audio(filepath, self.sr, self.max_duration)
        mel_spec = extract_mel_spectrogram(audio, self.sr, self.n_mels)
        mel_spec = torch.tensor(mel_spec, dtype=torch.float32)
        label = torch.tensor(label, dtype=torch.float32)
        return mel_spec, label

def get_file_paths_and_labels(dataset_dir):
    """
    从数据集目录加载音频文件路径及标签
    :param dataset_dir: 数据集根目录
    """
    audio_dir = os.path.join(dataset_dir, "audio_files")
    label_file = os.path.join(dataset_dir, "labels.csv")
    audio_files = [os.path.join(audio_dir, fname) for fname in os.listdir(audio_dir) if fname.endswith('.wav')]

    labels_df = pd.read_csv(label_file)
    labels = labels_df.set_index("filename").reindex([os.path.basename(f) for f in audio_files]).fillna(0).values
    return audio_files, labels
